package commands

import (
	"ConvertFlow/protocols/icmp"
	"ConvertFlow/utils"
	"github.com/sirupsen/logrus"
	"github.com/spf13/cobra"
	"os"
)

// ICMPFlags ICMP协议专用的标志
type ICMPFlags struct {
	RemoteAddr string
	ListenAddr string
	ChunkSize  int
	Delay      int
	NoReply    bool
	Key        string
}

var icmpFlags = &ICMPFlags{}

var IcmpRootCmd = &cobra.Command{
	Use:   "icmp",
	Short: "Commands using icmp protocol",
	Long:  `All commands that use icmp protocol for data transmission`,
}

// 发送命令
var IcmpSendCmd = &cobra.Command{
	Use:   "send [string to send]",
	Short: "Send data using ICMP protocol",
	Long:  `Send data to a remote receiver using ICMP protocol.`,
	Args:  cobra.MinimumNArgs(1),
	Run:   runSendByIcmp,
}

// 文件发送命令
var IcmpSendFileCmd = &cobra.Command{
	Use:   "file [path of the file]",
	Short: "Send file content using ICMP protocol",
	Long:  `Send the content of a file to a remote receiver using ICMP protocol.`,
	Args:  cobra.ExactArgs(1),
	Run:   runSendFileByIcmp,
}

func init() {
	// 添加ICMP标志
	addICMPFlags(IcmpSendCmd)

	//将Send命令添加到Root命令中
	IcmpRootCmd.AddCommand(IcmpSendCmd)
	// 将file命令添加为send的子命令
	IcmpSendCmd.AddCommand(IcmpSendFileCmd)
}

// addICMPFlags 添加ICMP专用标志
func addICMPFlags(cmd *cobra.Command) {
	cmd.PersistentFlags().StringVarP(&icmpFlags.RemoteAddr, "remote", "r", "", "ICMP remote receiver address (required)")
	cmd.PersistentFlags().StringVarP(&icmpFlags.ListenAddr, "listen", "l", "0.0.0.0", "ICMP listen address")
	cmd.PersistentFlags().IntVarP(&icmpFlags.ChunkSize, "size", "s", 65488, "ICMP packet size")
	cmd.PersistentFlags().IntVarP(&icmpFlags.Delay, "delay", "d", 4, "ICMP packet delay")
	cmd.PersistentFlags().BoolVarP(&icmpFlags.NoReply, "noreply", "N", false, "ICMP no reply mode")
	cmd.PersistentFlags().StringVarP(&icmpFlags.Key, "key", "k", "", "ICMP encryption key")
	cmd.MarkFlagRequired("remote")
}

// runSendByIcmp 处理发送字符串的逻辑
func runSendByIcmp(cmd *cobra.Command, args []string) {
	var data string
	if icmpFlags.Key != "" {
		pubKey := utils.Base64ToPublicKey(icmpFlags.Key)
		enc := utils.Base64EncryptWithPublicKey([]byte(args[0]), pubKey) //same msg, and key but different output? normal behaviour
		data = enc
	} else {
		data = args[0]
	}

	savePcap := false
	if icmpFlags.NoReply {
		icmp.SendNoReply(icmpFlags.ListenAddr, icmpFlags.RemoteAddr, icmpFlags.ChunkSize, icmpFlags.Delay, data, savePcap)
	} else {
		icmp.SendReply(icmpFlags.ListenAddr, icmpFlags.RemoteAddr, icmpFlags.ChunkSize, icmpFlags.Delay, data)
	}
}

// runSendFileByIcmp 处理发送文件的逻辑
func runSendFileByIcmp(cmd *cobra.Command, args []string) {
	var (
		data string
		err  error
	)

	//Retreive file content
	b, err := os.ReadFile(args[0])
	if err != nil {
		logrus.Fatal("os.ReadFile failed:", err)
	}

	if icmpFlags.Key != "" {
		pubKey := utils.Base64ToPublicKey(icmpFlags.Key)
		enc := utils.Base64EncryptWithPublicKey(b, pubKey) //same msg, and key but different output? normal behaviour
		data = enc
	} else {
		data = string(b)
	}

	//send it
	savePcap := false
	if icmpFlags.NoReply {
		err = icmp.SendNoReply(icmpFlags.ListenAddr, icmpFlags.RemoteAddr, icmpFlags.ChunkSize, icmpFlags.Delay, data, savePcap)
		if err != nil {
			logrus.Error("icmp.SendNoReply() function failed, err:%s", err.Error())
			return
		}
	} else {
		icmp.SendReply(icmpFlags.ListenAddr, icmpFlags.RemoteAddr, icmpFlags.ChunkSize, icmpFlags.Delay, data)
	}
}
